#include <iostream>
#include <string>
#include <vector>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <cstring>
#include <fcntl.h>
#include <algorithm>

using namespace std;

#define PORT 8888
#define BUF_SIZE 1024

vector<int> client_fds;
vector<string> nicknames;

void broadcast(int sender_fd, const string& msg) {
    for (int fd : client_fds) {
        if (fd != sender_fd) {
            send(fd, msg.c_str(), msg.size(), 0);
        }
    }
}

void remove_client(int fd) {
    auto it = find(client_fds.begin(), client_fds.end(), fd);
    if (it != client_fds.end()) {
        int index = distance(client_fds.begin(), it);
        close(fd);
        client_fds.erase(it);
        nicknames.erase(nicknames.begin() + index);
    }
}

int main() {
    int server_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (server_fd < 0) {
        perror("socket");
        return 1;
    }

    sockaddr_in addr {};
    addr.sin_family = AF_INET;
    addr.sin_port = htons(PORT);
    addr.sin_addr.s_addr = INADDR_ANY;

    if (bind(server_fd, (sockaddr*)&addr, sizeof(addr)) < 0) {
        perror("bind");
        return 1;
    }

    if (listen(server_fd, 10) < 0) {
        perror("listen");
        return 1;
    }

    fd_set all_set, read_set;
    FD_ZERO(&all_set);
    FD_SET(server_fd, &all_set);
    int max_fd = server_fd;

    cout << "Server listening on port " << PORT << "...\n";

    while (true) {
        read_set = all_set;

        if (select(max_fd + 1, &read_set, nullptr, nullptr, nullptr) < 0) {
            perror("select");
            break;
        }

        if (FD_ISSET(server_fd, &read_set)) {
            sockaddr_in cli_addr {};
            socklen_t len = sizeof(cli_addr);
            int client_fd = accept(server_fd, (sockaddr*)&cli_addr, &len);
            if (client_fd < 0) continue;

            char name_buf[32];
            int n = recv(client_fd, name_buf, sizeof(name_buf) - 1, 0);
            if (n <= 0) {
                close(client_fd);
                continue;
            }
            name_buf[n] = '\0';

            bool duplicate = false;
            for (const auto& name : nicknames) {
                if (name == name_buf) {
                    const char* err = "[Server] Nickname already in use.\n";
                    send(client_fd, err, strlen(err), 0);
                    close(client_fd);
                    duplicate = true;
                    break;
                }
            }
            if (duplicate) continue;

            client_fds.push_back(client_fd);
            nicknames.emplace_back(name_buf);
            FD_SET(client_fd, &all_set);
            max_fd = max(max_fd, client_fd);

            string join_msg = "[Server] " + nicknames.back() + " joined.\n";
            broadcast(client_fd, join_msg);
            cout << join_msg;
        }

        for (size_t i = 0; i < client_fds.size(); ++i) {
            int fd = client_fds[i];
            if (FD_ISSET(fd, &read_set)) {
                char buf[BUF_SIZE];
                int n = recv(fd, buf, BUF_SIZE - 1, 0);
                if (n <= 0) {
                    string left_msg = "[Server] " + nicknames[i] + " left.\n";
                    broadcast(fd, left_msg);
                    cout << left_msg;
                    FD_CLR(fd, &all_set);
                    remove_client(fd);
                    --i;
                } else {
                    buf[n] = '\0';
                    string msg(buf);

                    if (msg == "/quit\n") {
                        string left_msg = "[Server] " + nicknames[i] + " left.\n";
                        broadcast(fd, left_msg);
                        cout << left_msg;
                        FD_CLR(fd, &all_set);
                        remove_client(fd);
                        --i;
                    } else if (msg.rfind("/to ", 0) == 0) {
                        string content = msg.substr(4);
                        size_t space_pos = content.find(' ');

                        if (space_pos == string::npos) {
                            const char* err = "[Server] Usage: /to <nickname> <message>\n";
                            send(fd, err, strlen(err), 0);
                            continue;
                        }

                        string target_name = content.substr(0, space_pos);
                        string real_msg = content.substr(space_pos + 1);

                        int target_fd = -1;
                        for (size_t j = 0; j < nicknames.size(); ++j) {
                            if (nicknames[j] == target_name) {
                                target_fd = client_fds[j];
                                break;
                            }
                        }

                        if (target_fd == -1) {
                            string err = "[Server] No such user: " + target_name + "\n";
                            send(fd, err.c_str(), err.size(), 0);
                            continue;
                        }

                        string from_name = nicknames[i];
                        string final_msg = "[Private from " + from_name + "] " + real_msg;
                        send(target_fd, final_msg.c_str(), final_msg.size(), 0);
                    } else if (msg == "/who\n") {
                        string list = "[Online users]\n";
                        for (const auto& name : nicknames) {
                            list += "- " + name + "\n";
                        }
                        send(fd, list.c_str(), list.size(), 0);
                    } else {
                        string full_msg = "[" + nicknames[i] + "] " + msg;
                        broadcast(fd, full_msg);
                        cout << full_msg;
                    }
                }
            }
        }
    }

    close(server_fd);
    return 0;
}

